################################################################################
#
# find_image.py - Module for finding SDSS image(s) covering RA/Dec, with 
#                 additional functionality such as mosaicking fields, and
#                 checking coverage.
#
# Notes: This script also requires:
#       - wcstools: sky2xy must be in $PATH
#       - swarp: must be in $PATH
#       - imagemagick: convert must be in $PATH
#       - $BOSS_PHOTOOBJ must be set to DR8 tree location (setup tree dr8)
#       - postgres: /usr/local/psotgresql/lib must be in $LD_LIBRARY_PATH 
#       - spectradb: must "setup spectradb" so that $SPECTRADB_DIR is set
#
# ** IMPORTANT ** If you import this script to use various functions elsewhere,
#   you must have already made an sqlalchemy database connection!
################################################################################

"""
todo: 
"""

__author__ = 'Adrian Price-Whelan <adrn@nyu.edu>'

# Standard library dependencies
import os, sys
import re
import time
import cStringIO as StringIO
import tarfile
import warnings
import subprocess
from uuid import uuid4

# Third-party libraries
import pyfits as pf
import numpy as np
from sqlalchemy.sql.expression import func
from astLib import astWCS

# Internal imports
import mangle

try:
    from hooloovookit.string2radec import * # ????
except ImportError:
    print "Couldn't find hooloovookit\nDid you setup hooloovookit?"
    sys.exit(1)

if __name__ != "__main__":
    from spectradb.PhotofieldModelClasses import *
    import hooloovookit.DatabaseConnection
    db = hooloovookit.DatabaseConnection.DatabaseConnection()
    session = db.Session()
    
# This squelches "Warning: File may have been truncated: actual file length (11520)
#   is smaller than the expected size (12211200)"
warnings.simplefilter('ignore')

#
# Create singleton of field list
#
class SDSSFindImage(object):
    _fieldlist= dict()
    _tmpfiles = list()
    
    def __new__(cls, *args, **kwargs):
        """Only create the fieldlist once as a singleton"""
        if not cls._fieldlist.has_key(cls):
            cls._fieldlist[cls]= object.__new__(cls)
        return cls._fieldlist[cls]
    
    def coarsecover(self, ra, dec, rad=0.0):
        """Given a position and a radius, find all fields that could
        potentially overlap the specified circle on the sky."""
        mdist = 0.18+rad
        fields = session.query(FieldList).filter(func.q3c_radial_query(FieldList.ra, FieldList.dec, ra, dec, mdist)).\
                                          filter(FieldList.rerun != '157').all()
        return fields

    def coarsecover_primary(self, ra, dec, rad=0):
        """Given a position and a radius, find all fields that could
        potentially overlap the specified circle on the sky that are
        primary fields **somewhere**."""
        mdist= 0.18+rad
        fields = session.query(FieldList).join(Window).\
                         filter(func.q3c_radial_query(FieldList.ra, FieldList.dec, ra, dec, mdist)).\
                         filter(FieldList.rerun != '157').all()
        return fields
    
    def coarsecover_poly(self, ra, dec, rad=0):
        """Given a position and a radius, find all unified primary
        polygons that could potentially overlap the specified circle
        on the sky."""
        mdist= 0.18+rad
        windows = session.query(Window).join(FieldList).\
                          filter(func.q3c_radial_query(FieldList.ra, FieldList.dec, ra, dec, mdist)).\
                          filter(FieldList.rerun != '157').all()
        return windows
    
    def readpar(self, filename):
        fp=open(filename, "r")
        lines= fp.read()
        pars=dict()
        for line in lines.split("\n"):
            words= line.split()
            if(len(words)>1):
                keyword= words[0]
                value= words[1]
                pars[keyword]=value
        fp.close()
        return pars
    
    def writepar(self, filename, pars):
        fp=open(filename, "w")
        for key in pars.keys():
            fp.write(key+' '+pars[key]+'\n')
        fp.close()
        return
    
    def hdrfix(self, hdrfile):
        """Fix a corrected frame header to standard WCS.
        Takes a file name, and actually alters the file."""
        self._tmpfiles.append(hdrfile)
        hdu = pf.open(hdrfile, "update")

        # Fix astrometric header
        hdr = hdu[0].header
        if(hdr.has_key('PLTRAH')):
            rdstr = '%2.2i%2.2i%08.5f%s%2.2i%2.2i%08.5f' % (hdr['PLTRAH'], hdr['PLTRAM'],
                                                        hdr['PLTRAS'], hdr['PLTDECSN'],
                                                        hdr['PLTDECD'], hdr['PLTDECM'],
                                                        hdr['PLTDECS'])
                                                        
            (crval1, crval2) = string2radec(rdstr)
            crpix1= hdr['PPO3']/hdr['XPIXELSZ']-hdr['CRPIX1']+0.5
            crpix2= hdr['PPO6']/hdr['YPIXELSZ']-hdr['CRPIX2']+0.5
            scale1= -hdr['XPIXELSZ']/1000.
            scale2= hdr['YPIXELSZ']/1000.
            cd1_1= scale1*hdr['amdx1']/3600.
            cd1_2= scale2*hdr['amdx2']/3600.
            cd2_1= scale1*hdr['amdy2']/3600.
            cd2_2= scale2*hdr['amdy1']/3600.
            hdr['CTYPE1']='RA---TAN'
            hdr['CTYPE2']='DEC--TAN'
            hdr['CRPIX1']=crpix1
            hdr['CRPIX2']=crpix2
            hdr['CRVAL1']=crval1
            hdr['CRVAL2']=crval2
            hdr.update('CD1_1',cd1_1,'RA deg per column pixel')
            hdr.update('CD1_2',cd1_2,'RA deg per row pixel')
            hdr.update('CD2_1',cd2_1,'Dec deg per column pixel')
            hdr.update('CD2_2',cd2_2,'Dec deg per row pixel')
            
        try:
            # Fix astronomy HDU header
            hdr = hdu[3].header
            hdr.update('TTYPE14', 'F')
        except IndexError:
            # When called from gethdr, it only copies the first hdu
            pass
       
        # Fix everything else possible
        hdu.verify('fix')
        hdu.close()
        del hdu
        return
    
    def gethdr(self, field, filter='r'):
        """Pull out and correct a header for a field, returning a file
        name for the fixed header in /tmp."""
        fbase = field.fitsBase(filter)
        fname = field.fitsPath(filter)
        pbunzip = subprocess.Popen(['bunzip2', '-c', '-k', fname],
                                  stdout=subprocess.PIPE)
        phead = subprocess.Popen(['head', '-c', '11520'], stdin=pbunzip.stdout,
                                stdout=subprocess.PIPE)
        hdrfile = '/tmp/hdr-'+fbase+'-'+str(phead.pid)
        
        if hdrfile not in self._tmpfiles:
            self._tmpfiles.append(hdrfile)
            
        fp = open(hdrfile, 'w')
        fp.write((phead.communicate())[0])
        fp.close()
        pbunzip.kill()
        self.hdrfix(hdrfile)
        return hdrfile
    
    def radec2xy(self, field, ra, dec, filter='r'):
        """Call sky2xy for a field and return X and Y given an RA and DEC.
        Assumes the r-band by default."""
        hdrfile = self.gethdr(field, filter=filter)
        sra = "%20.10f" % ra
        sdec = "%20.10f" % dec
        xout =subprocess.Popen(['sky2xy', '-o', 'x', hdrfile, sra, sdec],\
                                stdout=subprocess.PIPE)
        x = float(xout.stdout.read())-1.
        yout =subprocess.Popen(['sky2xy', '-o', 'y', hdrfile, sra, sdec],\
                                stdout=subprocess.PIPE)
        y = float(yout.stdout.read())-1.
        
        # clean up mess in /tmp
        #if os.path.exists(hdrfile): os.remove(hdrfile)
            
        return (x, y)
    
    def clean(self):
        """ Deletes the tmp files that are kept track of by the class"""
        for tmpfile in self._tmpfiles:
            if os.path.exists(tmpfile):
                os.remove(tmpfile)
            self._tmpfiles.pop(self._tmpfiles.index(tmpfile))
    
    def primary(self, ra, dec):
        """Find the primary field covering a given position on the sky."""
        windows = self.coarsecover_poly(ra, dec, rad=1./3600.)        
        if len(windows) == 0: raise ValueError("No primary field found for given position RA: %.5f, Dec: %.5f" % (ra, dec))
        
        polys = mangle.Mangle(windows, db=True)
        pid = polys.polyid(np.array([ra]), np.array([dec]))
        
        field = session.query(FieldList).filter(FieldList.index==windows[pid].ifield).\
                                         filter(FieldList.rerun != '157').all()
        return field[0]
        
    def jpgClip(self, filename, field, x, y, writedir="/tmp", size=100, northup=False):
        """Clip out a piece of a JPG, given field, x, y and size"""
        (xmax, ymax) = (2047, 1488)
        border = size
        xst = border + (x - size/2) 
        yst = border + ((ymax-y) - size/2) 
        jpgname = field.jpgPath
        
        os.system('convert -bordercolor black -border '+str(border)+'x'+str(border)+
                  ' -crop '+str(size)+'x'+str(size)+'+'+str(xst)+'+'+str(yst)+' '+
                  jpgname+' '+filename)
                  
        if northup:
            fitsname = field.fitsPath("r")
            pbunzip = subprocess.Popen(['bunzip2', '-c', '-k', fitsname],
                                     stdout=subprocess.PIPE)
            fitsData = (pbunzip.communicate())[0]            
            uncname = os.path.basename(os.path.splitext(fitsname)[0])
            
            fp = open(os.path.join(writedir, uncname), "w")
            fp.write(fitsData)
            fp.close()
            
            self.hdrfix(os.path.join(writedir, uncname))
            
            fieldWCS = astWCS.WCS(os.path.join(writedir, uncname),0)
            rotAngle = fieldWCS.getRotationDeg()
       
            os.system("convert -rotate %.2f -background black " % rotAngle + filename + " " + filename)
        
        self._tmpfiles.append(filename)
        f = open(filename, 'r')
        fVal = f.read()
        f.close()
        self.clean()
        
        return fVal
    
    def jpgCutout(self, ra, dec, size, write=True, writedir="/tmp", northup=False):
        img = SDSSFindImage()
        field = img.primary(ra, dec)
        
        x,y = img.radec2xy(field, ra, dec)    
        filename = '/tmp/tmp_thumb-%s.jpg' % uuid4()
        img.jpgClip(filename, field, x, y, size=size, northup=northup, writedir=writedir)
        
        if write:
            return filename
        else:
            f = open(filename, 'r')
            image = StringIO.StringIO(f.read())
            f.close()
        
            os.remove(filename)
            bitStream = image.getvalue()
            image.close()
            return bitStream
    
    def mosaic(self, ra, dec, outputFilename, writedir, fields, filters="r", pixelscale=0.0, size=0.05):
        """swarp together a mosaic around given RA/Dec"""
        
        if not os.path.isdir(writedir):
            os.makedirs(writedir)
        
        mosaickedImageFiles = []
        outputFilePrefix, outputFileExtension = os.path.splitext(outputFilename)
        for filter in np.unique(list(filters)):
            thisFile = outputFilePrefix + "-%s.fits" % filter
            
            # unzip all fields into writedir and fix the FITS headers
            names=list()
            for field in fields:
                fitsname = field.fitsPath(filter)
                pbunzip=subprocess.Popen(['bunzip2', '-c', '-k', fitsname],
                                         stdout=subprocess.PIPE)

                uncname = os.path.basename(os.path.splitext(fitsname)[0])
                fp = open(os.path.join(writedir, uncname), "w")
                fp.write((pbunzip.communicate())[0])
                fp.close()
                
                self.hdrfix(os.path.join(writedir, uncname))
                names.append(uncname)
            
            # build default.swarp
            
            def_swarp = os.path.join(os.environ['SASWEBAPP_DIR'],'saswebapp','config','default.swarp')
            swarp_params = self.readpar(def_swarp)
            
            # the output filename for the final mosaicked image
            swarp_params['IMAGEOUT_NAME'] = thisFile
            
            # the center ra and dec of the final mosaic
            swarp_ra, swarp_dec = ("%20.10f" % ra, "%20.10f" % dec)
            swarp_params['CENTER'] = swarp_ra + ',' + swarp_dec
            
            # the size of the mosaick in pixels
            swarp_size = str(long(size*3600./0.396))
            swarp_params['IMAGE_SIZE'] = swarp_size + ',' + swarp_size
            
            # if the user specified a pixel size, set this parameter
            if pixelscale != 0.0: # pixel scale in arcseconds / pixel
                pixelscale = float(pixelscale)
                pixelscale_type = 'MANUAL'
                swarp_params['PIXEL_SCALE'] = pixelscale
                swarp_params['PIXELSCALE_TYPE'] = pixelscale_type
            
            # write these out to a new default.swarp
            self.writepar(os.path.join(writedir, 'default.swarp'), swarp_params)
            
            # run swarp
            cmd= "cd " + writedir + " ; swarp "
            for name in names:
                cmd = cmd + " " + name
            os.system(cmd)
            
            self._tmpfiles += names
            self.clean()
        
            mosaickedImageFiles.append(os.path.join(writedir, thisFile))
            
        return mosaickedImageFiles
    
    def mosaic2tar(self, ra, dec, filters="r", size=0.05, pixelscale=0.0, primary=True):
        """swarp together a mosaic around given RA/Dec, return a bitstream of the tarfile"""
        assert size > 0
        radius = size / 1.414 + 0.16
        
        writedir = "/tmp/saswebapp-swarp-%s" % uuid4()
        outfile = 'J%s.fits' % radec2string(ra,dec).replace(':','')
        
        if primary: fields = self.coarsecover_primary(ra,dec,radius)
        else: fields = self.coarsecover(ra,dec,radius)
        
        fitsList = self.mosaic(ra, dec, outfile, writedir, fields, filters=filters, pixelscale=pixelscale, size=size)
        
        tarStringIO = StringIO.StringIO()
        tar = tarfile.open(fileobj=tarStringIO, mode='w')
        
        toClose = []
        mtime = time.time()
        for fitsFile in fitsList:
            openFits = open(fitsFile, 'r')
            fitsData = openFits.read()
            openFits.close()
            self._tmpfiles.append(fitsFile)
            
            fitsFileStringIO = StringIO.StringIO(fitsData)
            toClose.append(fitsFileStringIO)
            
            tarinfo = tarfile.TarInfo(name=os.path.basename(fitsFile))
            tarinfo.size = fitsFileStringIO.len
            tarinfo.mtime = mtime
            tar.addfile(tarinfo, fileobj=fitsFileStringIO)
        
        self._tmpfiles.append(os.path.join(writedir,'default.swarp'))
        self._tmpfiles.append(os.path.join(writedir,'coadd.weight.fits'))
        self.clean()
        os.rmdir(writedir)
        
        for file in toClose: file.close()
        
        tarData = tf.getvalue()
        tf.close()
        
        return tarData
        
# Convenience functions
def RADec2RunCamcolField(ra, dec):
    img = SDSSFindImage()
    field = img.primary(ra, dec)
    return field.run, field.camcol, field.field

def RADec2FieldObj(ra, dec):
    img = SDSSFindImage()
    field = img.primary(ra, dec)
    return field

def RADec2ThumbStringIO(ra, dec, size=100):
    img = SDSSFindImage()
    field = img.primary(ra, dec)
    x,y = img.radec2xy(field, ra, dec)
    
    filename = '/tmp/tmp_thumb-%s.jpg' % uuid4()
    img.jpgClip(filename, field, x, y, size, northup=True)
    f = open(filename, 'r')
    image = StringIO.StringIO(f.read())
    f.close()
    os.remove(filename)
    tmp = image.getvalue()
    image.close()
    return tmp

def check_cover(run, camcol, fieldid):
    num = session.query(FieldList).filter(FieldList.run == run).\
                                filter(FieldList.camcol == camcol).\
                                filter(FieldList.field == fieldid).count()
    if num == 0: return False
    return True
    
if __name__ == "__main__":
    from optparse import OptionParser
    
    parser = OptionParser(description="Create the BOSS and SHARED schemas.")
    parser.add_option("-H", "--host", dest="host", default="localhost",
                    help="Specify the host name of the database server.")
    
    parser.add_option("-D", "--database", dest="database", default="spectradb",
                    help="Specify the database name.")
    
    parser.add_option("-P", "--port", dest="port", default=5432, type=int,
                    help="The port number of the database server.")
    
    parser.add_option("-U", "--user", dest="user", default="postgres",
                    help="Specify the user name of the database.")
    
    parser.add_option("-v", "--verbose", action="store_true", dest="verbose", default=False,
                    help="Be chatty.")
    
    (options, args) = parser.parse_args()
    dbDict = {"db" : options.database, "port" : options.port, "host" : options.host, "user" : options.user}
    
    from hooloovookit.DatabaseConnection import DatabaseConnection
    db = DatabaseConnection(host=dbDict['host'], user=dbDict['user'], port=dbDict['port'], name=dbDict['db'])
    from spectradb.PhotofieldModelClasses import *
    
    # DO TESTS
    session = db.Session
    
    # Test jpg rotation:
    ra,dec = (237.74594, 1.4850507)
    size = 100
    img = SDSSFindImage()
    field = img.primary(ra, dec)
    x,y = img.radec2xy(field, ra, dec)
    filename = '/tmp/tmp_thumb-%s.jpg' % uuid4()
    img.jpgClip(filename, field, x, y, size=size, northup=True, writedir="/tmp")
    print filename